#!/usr/bin/env ruby

# == Synopsis
# Analyze the maximum stack usage of a program
#
# == Usage
#   stack_anal.rb -h for detailed help
#
# == Author
# Ratish J. Punnoose
#
# == License
# License: GPLv2
# No warranties at all.

#------------ Required libraries ---------------------------
require 'optparse'
#-----------------------------------------------------------

# Add the directory containing this script to the load path.
$:.unshift File.dirname(__FILE__)
require 'fnanalyze'






# ---------------------------------------- Default values 
inputstr = $stdin;   #Input from stdin

fptropts = Hash.new   
opts = OptionParser.new
mfunc_name='main'      #starting function is main
showcallees = false  # don't show the called functions



## What architectures do we currently support.
## Get the list of architectures dynamically
disassembly_archs = Disassembler.singleton_methods(false).sort

# Get descriptions for all archs
arch_help_str = ''
disassembly_archs.each{|arch|
    arch_help_str += DescribeUsage.send(arch)
}

## Default is the first one.
default_arch = disassembly_archs[0].dup
disassembly_archs.each{ |k| k.insert(0, ' ') }
# ------------------------------------------------------------






# ------------------------------ Parse options --------------------

opts.on("-a arch",
        "Choose architecture of input disassembly: DEFAULT #{default_arch}",
        "OPTIONS: #{disassembly_archs}",
        "#{arch_help_str}\n",
        String) {|arch|
    if (Disassembler.respond_to?(arch))
        default_arch = arch
        puts DescribeUsage.send(default_arch)
    else
        puts "Architecture #{arch} is NOT supported"
        exit
    end
}

opts.on("-f inputfilename", 
        "unless input is coming from STDIN\n",
        String) {|file| 
    inputstr = File.open(file)
}

opts.on("-m mainfuncname",
        "name of the function to traverse and compute stack size\n",
        String){|x| mfunc_name  = x}

opts.on("-p fname=callee1[,callee2,calleeN]", 
        "where fname is the name of a function that has ",
        "one or more indirect function callees",  
        "(called via function pointers).",
        "callee1, callee2, ... are the functions that are",
        "called by fname indirectly",
        "This can be repeated as",
        "-p f1=c1,c2 -p f2=c3",
        String ) {|str| 
    tmp = str.split('=')
    if (tmp.length < 2)
        return
    end
    fname = tmp[0]
    # If there are multiple indirect callees, split them
    callees = tmp[1].split(",")
    fptropts[fname] = callees 
}


opts.on("-c",
        "Show all functions and callees\n") {showcallees = true}




opts.parse(ARGV)
# ------------------------------ Finished parsing options --------------------






# Analysze the disassembly
#------------------------------------------------------------
program = Disassembler.send(default_arch, inputstr)



# Fix all indirect calls (function pointers)
program.fixindcalls(fptropts)


# Analyze for stack usage, and recursion
program.each{ |fn|
    # Traverse each function
    # traverse requires the program to get a list of all functions 
    fn.traverse(program)
}
#  -------------------- Done Analysis










## Check for functions that call other functions via function pointers
funcs_with_indcallees = Array.new
program.each {|fn|
    if (fn.indirect_callees)
        funcs_with_indcallees.push(fn.to_s)
    end
}




puts "\n------------------------------------------------------------------"
# Print out stack usage of main function
mfunc = program.findbyname(mfunc_name);
if (!mfunc )
    puts "\n#{mfunc_name} not found"
else
    print "Max stack value for #{mfunc_name} is #{mfunc.stktotal}\n"
    mfunc.follow(program) { |f|
        print "#{f} (#{f.stklocal}/#{f.stktotal}) --> "
    }
    print("0\n")
end



if (not funcs_with_indcallees.empty?)
    print "\nFunctions with INDIRECT callees: " +
        funcs_with_indcallees.join(" ") + "\n"   
end




# -------------------- Do display
if (showcallees)
    puts "\n------------------------------------------------------------------"
    program.each {|fn|
        # print hex address, name, localusage
        print "\n" + "#{fn.hex_addr}:" + fn.to_s + " (#{fn.stklocal}/#{fn.stktotal} ):\t"
        fn.each {|callee_addr|
            callee = program.findbyaddr(callee_addr)
            if (callee)
                print "\t" + callee.to_s + "(#{callee.stklocal}/#{callee.stktotal})"
            else
                print "\tUNKNOWN"
            end
        }
        if (fn.indirect_callees)
            print "\tINDIRECT"
        end
    }
    print("\n");
end

